import os
from typing import TypedDict, Annotated

from IPython.display import Image, display
from PIL import Image as PILImage
from langchain_openai import ChatOpenAI
from langgraph.constants import START, END
from langgraph.graph import add_messages, StateGraph



class State(TypedDict):
    messages: Annotated[list,add_messages]

#定义图结构
graph_builder = StateGraph(State)

# 使用DeepSeek-chat模型作为chatmodel
# model = ChatOpenAI(model="deepseek-chat",
#                    api_key=os.environ.get("DEEPSEEK_API_KEY"),
#                    base_url=os.environ.get("DEEPSEEK_BASE_URL"))

model = ChatOpenAI(model="deepseek-chat",
                   api_key=os.environ.get("DEEPSEEK_API_KEY"),
                   base_url=os.environ.get("DEEPSEEK_BASE_URL"))

def chatbot(state: State):
    return {"messages": model.invoke(state["messages"])}

def buildGraph(graph_builder):
    graph_builder.add_node("chatbot", chatbot)
    graph_builder.add_edge(START, "chatbot")
    graph_builder.add_edge("chatbot", END)

    return graph_builder.compile()

graph = buildGraph(graph_builder)

# 假设 graph.get_graph(xray=True).draw_mermaid_png() 返回的是 PNG 图像的字节数据
image_data = graph.get_graph(xray=True).draw_mermaid_png()

# 将字节数据保存为临时文件再显示
with open('temp_image.png', 'wb') as f:
    f.write(image_data)

display(Image(filename='temp_image.png'))
# 加载图像
image = PILImage.open('temp_image.png').convert('L')  # 转换为灰度图像

# 定义字符集
ASCII_CHARS = "@%#*+=-:. "

def resize_image(image, new_width=100):
    width, height = image.size
    ratio = height / width / 1.65  # 调整比例以适应控制台字符
    new_height = int(new_width * ratio)
    resized_image = image.resize((new_width, new_height))
    return resized_image

def grayify(image):
    grayscale_image = image.convert("L")
    return grayscale_image

def pixels_to_ascii(image):
    pixels = image.getdata()
    ascii_str = "".join([ASCII_CHARS[pixel // 32] for pixel in pixels])
    img_width = image.width
    ascii_str_len = len(ascii_str)
    ascii_img = "\n".join([ascii_str[index:(index + img_width)] for index in range(0, ascii_str_len, img_width)])
    return ascii_img

# 调整图像大小
resized_image = resize_image(image)

# 将图像转换为灰度图像
gray_image = grayify(resized_image)

# 将像素转换为ASCII字符
ascii_str = pixels_to_ascii(gray_image)

# 打印ASCII字符图像
print(ascii_str)

#推理任务执行
# result = graph.invoke({"messages":"你好，请介绍一下你自己"})
# print(result["messages"][-1].content)

def stream_graph_updates(user_input: str):
    # for event in graph.invoke({"messages": [("user", user_input)]}):
    #     for value in event.values():
    #         print("模型回复:", value["messages"].content)
    result = graph.invoke({"messages": [("user", user_input)]})
    print("模型回复:", result["messages"][-1].content)

#多轮对话任务执行
while True:
    try:
        user_input = input("用户提问: ")
        if user_input.lower() in ["退出"]:
            print("下次再见！")
            break

        stream_graph_updates(user_input)
    except:
        break
